package server

import (
	"context"
	"sync"

	"gitee.com/zfd81/dbrs/log"

	"gitee.com/zfd81/dbrs/errs"

	"gitee.com/zfd81/dbrs/engine"

	sys2 "gitee.com/zfd81/dbrs/sys"
	"github.com/opentracing/opentracing-go"
	"github.com/sirupsen/logrus"
	"vitess.io/vitess/go/mysql"
)

// SessionBuilder creates sessions given a MySQL connection and a server address.
type SessionBuilder func(ctx context.Context, conn *mysql.Conn, addr string) (engine.Session, error)

// DoneFunc is a function that must be executed when the session is used and
// it can be disposed.
type DoneFunc func()

// DefaultSessionBuilder is a SessionBuilder that returns a base session.
func DefaultSessionBuilder(ctx context.Context, c *mysql.Conn, addr string) (engine.Session, error) {
	client := engine.Client{Address: c.RemoteAddr().String(), User: c.User, Capabilities: c.Capabilities}
	ps, err := initProviders()
	if err != nil {
		return nil, err
	}
	sess := engine.NewSession(addr, client, c.ConnectionID, ps)
	db := conf.Databases[0]
	user, err := sys2.FindUser(c.User)
	if err != nil {
		return nil, err
	}
	if user != nil && user.DefaultSignature != "" {
		sess.SetUserVariable("user_sign", user.DefaultSignature)
	}
	sess.SetCurrentDatabase(db.Name)
	return sess, nil
}

// SessionManager is in charge of creating new sessions for the given
// connections and keep track of which sessions are in each connection, so
// they can be cancelled if the connection is closed.
type SessionManager struct {
	addr      string
	tracer    opentracing.Tracer
	hasDBFunc func(name string) bool
	mu        *sync.Mutex
	builder   SessionBuilder
	sessions  map[uint32]engine.Session
	pid       uint64
}

// NewSessionManager creates a SessionManager with the given SessionBuilder.
func NewSessionManager(
	builder SessionBuilder,
	tracer opentracing.Tracer,
	hasDBFunc func(name string) bool,
	addr string,
) *SessionManager {
	return &SessionManager{
		addr:      addr,
		tracer:    tracer,
		hasDBFunc: hasDBFunc,
		mu:        new(sync.Mutex),
		builder:   builder,
		sessions:  make(map[uint32]engine.Session),
	}
}

func (s *SessionManager) nextPid() uint64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pid++
	return s.pid
}

// NewSession creates a Session for the given connection and saves it to the session pool.
func (s *SessionManager) NewSession(ctx context.Context, conn *mysql.Conn) error {
	var err error

	s.mu.Lock()
	defer s.mu.Unlock()
	s.sessions[conn.ConnectionID], err = s.builder(ctx, conn, s.addr)

	if err != nil {
		return err
	}

	logger := s.sessions[conn.ConnectionID].GetLogger()
	if logger == nil {
		log := logrus.StandardLogger()
		logger = logrus.NewEntry(log)
	}

	s.sessions[conn.ConnectionID].SetLogger(
		logger.WithField(log.ConnectionIdLogField, conn.ConnectionID).
			WithField(log.ConnectTimeLogKey, s.sessions[conn.ConnectionID].CreateTime().Format("2006-01-02 15:04:05.000000")),
	)
	sys2.ThreadsCreatedIncrease()
	return err
}

func (s *SessionManager) SetDB(conn *mysql.Conn, db string) error {
	sess, err := s.getOrCreateSession(context.Background(), conn)
	if err != nil {
		return err
	}

	if db != "" && !sess.HasDatabase(db) {
		return errs.ErrDatabaseNotFound.New(db)
	}

	sess.SetCurrentDatabase(db)
	return nil
}

func (s *SessionManager) session(conn *mysql.Conn) engine.Session {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sessions[conn.ConnectionID]
}

// NewContext creates a new context for the session at the given conn.
func (s *SessionManager) NewContext(conn *mysql.Conn) (*engine.Context, error) {
	return s.NewContextWithQuery(conn, "")
}

func (s *SessionManager) getSession(conn *mysql.Conn) engine.Session {
	return s.sessions[conn.ConnectionID]
}
func (s *SessionManager) getOrCreateSession(ctx context.Context, conn *mysql.Conn) (engine.Session, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	sess, ok := s.sessions[conn.ConnectionID]

	if !ok {
		s.mu.Unlock()
		err := s.NewSession(ctx, conn)
		s.mu.Lock()
		if err != nil {
			return nil, err
		}
		sess = s.sessions[conn.ConnectionID]
	}

	return sess, nil
}

// NewContextWithQuery creates a new context for the session at the given conn.
func (s *SessionManager) NewContextWithQuery(conn *mysql.Conn, query string) (*engine.Context, error) {
	ctx := context.Background()
	sess, err := s.getOrCreateSession(ctx, conn)

	if err != nil {
		return nil, err
	}

	context := engine.NewContext(
		ctx,
		engine.WithSession(sess),
		engine.WithTracer(s.tracer),
		engine.WithPid(s.nextPid()),
		engine.WithQuery(query),
		engine.WithRootSpan(s.tracer.StartSpan("query")),
	)

	return context, nil
}

// CloseConn closes the connection in the session manager and all its
// associated contexts, which are cancelled.
func (s *SessionManager) CloseConn(conn *mysql.Conn) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.sessions, conn.ConnectionID)
}

func initProviders() (map[string]engine.Provider, error) {
	ps := map[string]engine.Provider{}
	for _, db := range conf.Databases {
		provider, err := engine.NewProvider(db)
		if err != nil {
			return nil, err
		}
		ps[provider.Name()] = provider
	}
	return ps, nil
}
